#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>
#include <sched.h>
#include <pthread.h>

#include "common.h"
#include "msr_address.h"
#include "cpu_info.h"
#include "generate_code.h"

uint64_t g_result[10];
CodeConfig *g_CodeCfg;

void *g_code;
size_t g_code_len;
size_t g_code_size;

uint64_t *g_stack;

uint8_t *_measure = NULL;
size_t _measure_len = 0;

size_t mmap_bin(const char *file_name, uint8_t **des)
{
    int fd = open(file_name, O_RDONLY);
    size_t len = lseek(fd, 0, SEEK_END);
    *des = mmap(NULL, len, PROT_READ, MAP_PRIVATE, fd, 0);
    if (*des == MAP_FAILED) {
        printf("Error map file: %s\n", file_name);
        len = -1;
    }
    close(fd);
    return len;
}

int read_msr(int cpu, uint32_t address, uint64_t *val)
{
    char cmd[50];
    snprintf(cmd, sizeof(cmd), "rdmsr -c -p%d %#x", cpu, address);
    FILE *fp;
    if (!(fp = popen(cmd, "r"))) {
        printf("Error reading from \"%s\"", cmd);
        return -1;
    }

    char buf[20];
    if (NULL == fgets(buf, sizeof(buf), fp)) {
        buf[0] = 0;
    }
    pclose(fp);

    *val = strtoul(buf, NULL, 0);
    return 0;
}

int write_msr(int cpu, uint32_t address, uint64_t value)
{
    char cmd[50];
    snprintf(cmd, sizeof(cmd), "wrmsr -p%d %#x %#lx", cpu, address, value);
    if (system(cmd)) {
        printf("\"%s\" failed. You may need to disable \"Secure Boot\" in BIOS.\n", cmd);
        return -1;
    }
    return 0;
}

int configure_fixed_pmu(int cpu)
{
    uint64_t global_ctrl;
    if (read_msr(cpu, MSR_IA32_PERF_GLOBAL_CTRL, &global_ctrl)) {
        return -1;
    }
    uint64_t fixed_ctrl;
    if (read_msr(cpu, MSR_IA32_FIXED_CTR_CTRL, &fixed_ctrl)) {
        return -1;
    }

    global_ctrl |= ((uint64_t)7 << 32) | 15;
    if (write_msr(cpu, MSR_IA32_PERF_GLOBAL_CTRL, global_ctrl)) {
        return -1;
    }

    // disable fixed counters
    fixed_ctrl &= ~((1 << 12) - 1);
    if (write_msr(cpu, MSR_IA32_FIXED_CTR_CTRL, fixed_ctrl)) {
        return -1;
    }
    // clear
    for (int i = 0; i < 3; i++) {
        if (write_msr(cpu, MSR_IA32_FIXED_CTR0 + i, 0)) {
            return -1;
        }
    }
    // enable fixed counters
    uint32_t os = 0, usr = 1;
    fixed_ctrl |= (os << 8) | (os << 4) | os;
    fixed_ctrl |= (usr << 9) | (usr << 5) | (usr << 1);
    if (write_msr(cpu, MSR_IA32_FIXED_CTR_CTRL, fixed_ctrl)) {
        return -1;
    }
    return 0;
}

int configure_programmable_pmu(int cpu, PmuDataList *pmu_events, int pmu_events_cnt, int pmu_counters)
{
    uint64_t global_ctrl;
    if (read_msr(cpu, MSR_IA32_PERF_GLOBAL_CTRL, &global_ctrl)) {
        return -1;
    }
    global_ctrl |= ((uint64_t)7 << 32) | 15;
    if (write_msr(cpu, MSR_IA32_PERF_GLOBAL_CTRL, global_ctrl)) {
        return -1;
    }

    for (int i = 0; i < pmu_counters && i < pmu_events_cnt; i++) {
        if (pmu_events[i].data < FIX_PMU_EVENT_MAX) {
            continue;
        }

        uint64_t pmu_event;
        if (read_msr(cpu, MSR_IA32_PERFEVTSEL0 + i, &pmu_event)) {
            return -1;
        }
        // disable counter i
        pmu_event &= ~(((uint64_t)1 << 32) - 1);
        write_msr(cpu, MSR_IA32_PERFEVTSEL0 + i, pmu_event);
        // clear
        write_msr(cpu, MSR_IA32_PMC0 + i, 0);
        // configure counter i
        pmu_event = pmu_events[i].data;
        write_msr(cpu, MSR_IA32_PERFEVTSEL0 + i, pmu_event);
    }
    return 0;
}

void generate_code(int loop, int unroll, PmuDataList *events, int e_cnt, int pmu_counters)
{
    uint64_t event;
    uint8_t *gen = g_code;
    size_t gen_len = 0;
    int loop_offset;

    gen_len += reg_push(gen + gen_len);

    // mov r15, g_result
    gen_len += mov_reg_imm64(gen + gen_len, R15, (uint64_t)g_result);

    // read init pmu data
    for (int i = 0; i < pmu_counters && i < e_cnt; i++) {
        gen_len += lfence(gen + gen_len);
        event = events[i].data;
        if (event == RDTSC) {
            gen_len += rdtsc(gen + gen_len);
        } else {
            if (event == INSTRUCTIONS) {
                gen_len += mov_reg_imm32(gen + gen_len, RCX, 0x40000000);
            } else if (event == CYCLES) {
                gen_len += mov_reg_imm32(gen + gen_len, RCX, 0x40000001);
            } else {
                gen_len += mov_reg_imm32(gen + gen_len, RCX, i);
            }
            gen_len += lfence(gen + gen_len);
            gen_len += rdpmc(gen + gen_len);
        }
        gen_len += lfence(gen + gen_len);
        gen_len += sub_rdx_rax_to_r15(gen + gen_len, i * 8);
    }

    if (loop > 0) {
        // mov rdi, loop
        gen_len += mov_reg_imm32(gen + gen_len, RDI, loop);

        // aligned 0x40
        size_t o = 0x40 - (size_t)(gen_len & (0x40 - 1));
        while (o > 0) {
            // nop
            gen[gen_len++] = '\x90';
            o--;
        }
        loop_offset = gen_len;
    }

    // copy measure code
    memcpy(gen + gen_len, _measure, _measure_len);
    gen_len += _measure_len;

    // unroll code
    int add_unroll = unroll - 1;
    while (add_unroll > 0) {
        add_unroll--;
        memcpy(gen + gen_len, _measure, _measure_len);
        gen_len += _measure_len;
    }

    if (loop > 0) {
        // dec rdi
        gen_len += dec_reg(gen + gen_len, RDI);
        // jnz
        gen_len += jnz_imm32(gen + gen_len, loop_offset - gen_len);
    }

    // read fini pmu data
    for (int i = 0; i < pmu_counters && i < e_cnt; i++) {
        gen_len += lfence(gen + gen_len);
        event = events[i].data;
        if (event == RDTSC) {
            gen_len += rdtsc(gen + gen_len);
        } else {
            if (event == INSTRUCTIONS) {
                gen_len += mov_reg_imm32(gen + gen_len, RCX, 0x40000000);
            } else if (event == CYCLES) {
                gen_len += mov_reg_imm32(gen + gen_len, RCX, 0x40000001);
            } else {
                gen_len += mov_reg_imm32(gen + gen_len, RCX, i);
            }
            gen_len += lfence(gen + gen_len);
            gen_len += rdpmc(gen + gen_len);
        }
        gen_len += lfence(gen + gen_len);
        gen_len += add_rdx_rax_to_r15(gen + gen_len, i * 8);
    }

    gen_len += reg_pop(gen + gen_len);

    // ret
    gen[gen_len++] = '\xc3';
    g_code_len = gen_len;
    return;
}

void write_file(const char *file_name)
{
    uint8_t *p = g_code;
    FILE *fp = fopen(file_name, "wb");
    if (NULL == fp) {
        printf("Error open file: %s\n", file_name);
        return;
    }
    fwrite(p, 1, g_code_len, fp);
    fclose(fp);
    return;
}

int bind_to_cpu(int *bind_cpu)
{
    if (-1 == *bind_cpu) {
        *bind_cpu = sched_getcpu();
    }
    cpu_set_t mask;
    CPU_ZERO(&mask);
    CPU_SET(*bind_cpu, &mask);

    if (sched_setaffinity(0, sizeof(cpu_set_t), &mask) == -1) {
        printf("Error: Could not pin thread to core %d\n", *bind_cpu);
        return -1;
    }
    return 0;
}

int init(CodeConfig *codeCfg)
{
    g_CodeCfg = codeCfg;
    if (0 != bind_to_cpu(&g_CodeCfg->core)) {
        return -1;
    }

    _measure_len = mmap_bin("_measure.bin", &_measure);

    // calloc memory to store measure code
    long int pg_size = sysconf(_SC_PAGESIZE);
    g_code_size = 900 + _measure_len * codeCfg->unroll;
    if (posix_memalign(&g_code, pg_size, g_code_size)) {
        printf("Error malloc failed!\n");
        return -1;
    }
    if (mprotect(g_code, g_code_size, PROT_READ | PROT_WRITE | PROT_EXEC)) {
        printf("Error mprotect failed!\n");
        return -1;
    }
    return 0;
}

void measure(CodeConfig *codeCfg)
{
    int pmu_counters = codeCfg->pmu_cnt;
    PmuDataList *events = codeCfg->events_list;
    int e_cnt = codeCfg->events_cnt;
    int run_cnt = codeCfg->run;

    // enable fixed pmu
    if (configure_fixed_pmu(g_CodeCfg->core)) {
        return;
    }

    for (int i = 0; i < e_cnt; i += pmu_counters) {
        memset(g_result, 0, sizeof(g_result));
        // enable programmable pmu
        if (configure_programmable_pmu(codeCfg->core, events + i, e_cnt - i, pmu_counters)) {
            return;
        }
        // generate measure g_code
        generate_code(codeCfg->loop, codeCfg->unroll, events + i, e_cnt - i, pmu_counters);

        // run
        for (int r = 0; r < run_cnt; r++) {
            ((void (*)())g_code)();
        }

        for (int c = i; c < i + pmu_counters && c < e_cnt; c++) {
            codeCfg->events_list[c].value = g_result[c] / run_cnt;
        }
    }

    if (NULL != codeCfg->binary_file) {
        write_file(codeCfg->binary_file);
    }

    if (NULL != g_code) {
        free(g_code);
    }

    return;
}

void start(CodeConfig *codeCfg)
{
    if (init(codeCfg)) {
        return;
    }
    measure(codeCfg);
    return;
}

int g_get_freq_running = 1;
PmuData g_get_freq_t1 = {0};
PmuData g_get_freq_t2 = {0};

void *test(void *arg)
{
    READ_CORE_CYCLE(g_get_freq_t1);
    while (g_get_freq_running) {
        READ_CORE_CYCLE(g_get_freq_t2);
    }
    return NULL;
}

void get_run_frequency(CodeConfig *codeCfg)
{
    if (0 != bind_to_cpu(&codeCfg->core)) {
        return;
    }
    if (configure_fixed_pmu(codeCfg->core)) {
        return;
    }

    pthread_t tid;
    if (0 != pthread_create(&tid, NULL, test, NULL)) {                                                                                          
        printf("error !\n");
    }

    g_get_freq_running = 1;
    sleep(5);
    g_get_freq_running = 0;
    pthread_join(tid, NULL);
    printf("running frequency : %.2f MHz\n", (g_get_freq_t2.t64 - g_get_freq_t1.t64) / 5 / 1e6);
    return;
}
